package com.mm.cloth.rmbg;

import android.content.Context;
import android.content.res.AssetManager;
import android.graphics.Bitmap;
import android.graphics.Color;
import android.util.Log;

import org.tensorflow.lite.Interpreter;
import org.tensorflow.lite.gpu.CompatibilityList;
import org.tensorflow.lite.gpu.GpuDelegate;

import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.MappedByteBuffer;
import java.nio.channels.FileChannel;

/**
 * 基于TensorFlow Lite的Watermark-Removal推理助手
 * 专门用于水印检测和去除
 */
public class WatermarkRemovalTFLiteHelper {
    private static final String TAG = "WatermarkRemovalTFLite";
    
    private Interpreter interpreter;
    private GpuDelegate gpuDelegate;
    private boolean useGPU = false;
    
    // 模型参数（根据参考文档）
    private static final int MODEL_INPUT_SIZE = 256;  // 256x256输入尺寸
    private static final int CHANNELS = 3;            // RGB通道
    private static final int BATCH_SIZE = 1;
    
    // 输入输出缓冲区
    private ByteBuffer inputBuffer;
    private ByteBuffer outputBuffer;
    
    // 模型文件名
    private static final String MODEL_FILENAME = "watermark_removal.tflite";
    
    /**
     * 初始化TensorFlow Lite模型
     * @param context 应用上下文
     */
    public void initModel(Context context) throws Exception {
        Log.d(TAG, "初始化Watermark-Removal TensorFlow Lite模型");
        
        try {
            // 检查GPU支持
            CompatibilityList compatList = new CompatibilityList();
            if (compatList.isDelegateSupportedOnThisDevice()) {
                Log.d(TAG, "设备支持GPU加速");
                useGPU = true;
            } else {
                Log.d(TAG, "设备不支持GPU加速，使用CPU");
                useGPU = false;
            }
            
            // 加载模型文件
            MappedByteBuffer modelBuffer = loadModelFile(context.getAssets());
            
            // 配置解释器选项
            Interpreter.Options options = new Interpreter.Options();
            options.setNumThreads(4); // 设置CPU线程数
            
            if (useGPU) {
                try {
                    gpuDelegate = new GpuDelegate();
                    options.addDelegate(gpuDelegate);
                    Log.d(TAG, "GPU代理添加成功");
                } catch (Exception e) {
                    Log.w(TAG, "GPU代理创建失败，回退到CPU", e);
                    useGPU = false;
                    gpuDelegate = null;
                }
            }
            
            // 创建解释器
            interpreter = new Interpreter(modelBuffer, options);
            
            // 初始化输入输出缓冲区
            initBuffers();
            
            Log.d(TAG, "Watermark-Removal TensorFlow Lite模型初始化成功");
            Log.d(TAG, "使用" + (useGPU ? "GPU" : "CPU") + "进行推理");
            
            // 打印模型信息
            printModelInfo();
            
        } catch (Exception e) {
            Log.e(TAG, "Watermark-Removal TensorFlow Lite模型初始化失败", e);
            release();
            throw e;
        }
    }
    
    /**
     * 加载模型文件
     */
    private MappedByteBuffer loadModelFile(AssetManager assetManager) throws IOException {
        try (InputStream inputStream = assetManager.open(MODEL_FILENAME);
             FileInputStream fileInputStream = (FileInputStream) inputStream;
             FileChannel fileChannel = fileInputStream.getChannel()) {
            
            long startOffset = 0;
            long declaredLength = fileChannel.size();
            
            Log.d(TAG, "模型文件大小: " + (declaredLength / (1024 * 1024)) + " MB");
            
            return fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength);
        }
    }
    
    /**
     * 初始化输入输出缓冲区
     */
    private void initBuffers() {
        // 输入缓冲区: BATCH_SIZE * HEIGHT * WIDTH * CHANNELS * 4 bytes (float32)
        int inputSize = BATCH_SIZE * MODEL_INPUT_SIZE * MODEL_INPUT_SIZE * CHANNELS * 4;
        inputBuffer = ByteBuffer.allocateDirect(inputSize);
        inputBuffer.order(ByteOrder.nativeOrder());
        
        // 输出缓冲区: 与输入相同尺寸
        int outputSize = BATCH_SIZE * MODEL_INPUT_SIZE * MODEL_INPUT_SIZE * CHANNELS * 4;
        outputBuffer = ByteBuffer.allocateDirect(outputSize);
        outputBuffer.order(ByteOrder.nativeOrder());
        
        Log.d(TAG, "缓冲区初始化完成 - 输入: " + inputSize + " bytes, 输出: " + outputSize + " bytes");
    }
    
    /**
     * 打印模型信息
     */
    private void printModelInfo() {
        try {
            Log.d(TAG, "=== Watermark-Removal TFLite模型信息 ===");
            Log.d(TAG, "输入张量数量: " + interpreter.getInputTensorCount());
            Log.d(TAG, "输出张量数量: " + interpreter.getOutputTensorCount());
            
            // 输入张量信息
            for (int i = 0; i < interpreter.getInputTensorCount(); i++) {
                int[] inputShape = interpreter.getInputTensor(i).shape();
                Log.d(TAG, "输入张量 " + i + " 形状: [" + 
                      inputShape[0] + ", " + inputShape[1] + ", " + inputShape[2] + ", " + inputShape[3] + "]");
            }
            
            // 输出张量信息
            for (int i = 0; i < interpreter.getOutputTensorCount(); i++) {
                int[] outputShape = interpreter.getOutputTensor(i).shape();
                Log.d(TAG, "输出张量 " + i + " 形状: [" + 
                      outputShape[0] + ", " + outputShape[1] + ", " + outputShape[2] + ", " + outputShape[3] + "]");
            }
            
        } catch (Exception e) {
            Log.w(TAG, "获取模型信息失败", e);
        }
    }
    
    /**
     * 执行水印去除推理
     * @param inputBitmap 输入图像
     * @return 去水印后的图像
     */
    public Bitmap removeWatermark(Bitmap inputBitmap) throws Exception {
        if (interpreter == null) {
            throw new IllegalStateException("模型未初始化");
        }
        
        if (inputBitmap == null) {
            throw new IllegalArgumentException("输入图像不能为空");
        }
        
        Log.d(TAG, "开始Watermark-Removal推理");
        Log.d(TAG, "输入图像尺寸: " + inputBitmap.getWidth() + "x" + inputBitmap.getHeight());
        
        long startTime = System.currentTimeMillis();
        
        try {
            // 1. 预处理输入图像
            long preprocessStart = System.currentTimeMillis();
            preprocessInput(inputBitmap);
            long preprocessEnd = System.currentTimeMillis();
            
            // 2. 执行推理
            long inferenceStart = System.currentTimeMillis();
            interpreter.run(inputBuffer, outputBuffer);
            long inferenceEnd = System.currentTimeMillis();
            
            // 3. 后处理输出
            long postprocessStart = System.currentTimeMillis();
            Bitmap resultBitmap = postprocessOutput(inputBitmap.getWidth(), inputBitmap.getHeight());
            long postprocessEnd = System.currentTimeMillis();
            
            long totalTime = System.currentTimeMillis() - startTime;
            
            Log.d(TAG, "Watermark-Removal推理完成");
            Log.d(TAG, "时间统计 - 预处理: " + (preprocessEnd - preprocessStart) + "ms, " +
                       "推理: " + (inferenceEnd - inferenceStart) + "ms, " +
                       "后处理: " + (postprocessEnd - postprocessStart) + "ms, " +
                       "总计: " + totalTime + "ms");
            
            return resultBitmap;
            
        } catch (Exception e) {
            Log.e(TAG, "Watermark-Removal推理失败", e);
            throw e;
        }
    }
    
    /**
     * 预处理输入图像
     * 根据参考文档：归一化到[-1,1]范围
     */
    private void preprocessInput(Bitmap bitmap) {
        // 调整图像尺寸到模型输入尺寸
        Bitmap resizedBitmap = Bitmap.createScaledBitmap(bitmap, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE, true);
        
        // 获取像素数据
        int[] pixels = new int[MODEL_INPUT_SIZE * MODEL_INPUT_SIZE];
        resizedBitmap.getPixels(pixels, 0, MODEL_INPUT_SIZE, 0, 0, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE);
        
        // 清空输入缓冲区
        inputBuffer.rewind();
        
        // 转换像素数据为float并归一化到[-1,1]
        for (int pixel : pixels) {
            // 提取RGB值并归一化到[-1,1]
            float r = (Color.red(pixel) / 255.0f) * 2.0f - 1.0f;
            float g = (Color.green(pixel) / 255.0f) * 2.0f - 1.0f;
            float b = (Color.blue(pixel) / 255.0f) * 2.0f - 1.0f;
            
            // 按RGB顺序写入缓冲区
            inputBuffer.putFloat(r);
            inputBuffer.putFloat(g);
            inputBuffer.putFloat(b);
        }
        
        // 清理临时bitmap
        if (resizedBitmap != bitmap) {
            resizedBitmap.recycle();
        }
        
        Log.d(TAG, "输入预处理完成");
    }
    
    /**
     * 后处理输出数据
     * 将模型输出转换为Bitmap
     */
    private Bitmap postprocessOutput(int targetWidth, int targetHeight) {
        // 创建输出bitmap
        Bitmap outputBitmap = Bitmap.createBitmap(MODEL_INPUT_SIZE, MODEL_INPUT_SIZE, Bitmap.Config.ARGB_8888);
        
        // 重置输出缓冲区位置
        outputBuffer.rewind();
        
        // 读取输出数据并转换为像素
        int[] pixels = new int[MODEL_INPUT_SIZE * MODEL_INPUT_SIZE];
        
        for (int i = 0; i < pixels.length; i++) {
            // 读取RGB值（模型输出范围[-1,1]）
            float r = outputBuffer.getFloat();
            float g = outputBuffer.getFloat();
            float b = outputBuffer.getFloat();
            
            // 反归一化到[0,255]
            int red = Math.max(0, Math.min(255, (int) ((r + 1.0f) * 127.5f)));
            int green = Math.max(0, Math.min(255, (int) ((g + 1.0f) * 127.5f)));
            int blue = Math.max(0, Math.min(255, (int) ((b + 1.0f) * 127.5f)));
            
            pixels[i] = Color.argb(255, red, green, blue);
        }
        
        // 设置像素到bitmap
        outputBitmap.setPixels(pixels, 0, MODEL_INPUT_SIZE, 0, 0, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE);
        
        // 调整到目标尺寸
        Bitmap resultBitmap;
        if (targetWidth != MODEL_INPUT_SIZE || targetHeight != MODEL_INPUT_SIZE) {
            resultBitmap = Bitmap.createScaledBitmap(outputBitmap, targetWidth, targetHeight, true);
            outputBitmap.recycle();
        } else {
            resultBitmap = outputBitmap;
        }
        
        Log.d(TAG, "输出后处理完成，结果尺寸: " + resultBitmap.getWidth() + "x" + resultBitmap.getHeight());
        
        return resultBitmap;
    }
    
    /**
     * 释放资源
     */
    public void release() {
        try {
            if (interpreter != null) {
                interpreter.close();
                interpreter = null;
            }
            
            if (gpuDelegate != null) {
                gpuDelegate.close();
                gpuDelegate = null;
            }
            
            Log.d(TAG, "Watermark-Removal TensorFlow Lite资源已释放");
        } catch (Exception e) {
            Log.e(TAG, "释放Watermark-Removal TensorFlow Lite资源失败", e);
        }
    }
    
    /**
     * 检查是否已初始化
     */
    public boolean isInitialized() {
        return interpreter != null;
    }
    
    /**
     * 获取模型输入尺寸
     */
    public int getModelInputSize() {
        return MODEL_INPUT_SIZE;
    }
    
    /**
     * 是否使用GPU加速
     */
    public boolean isUsingGPU() {
        return useGPU;
    }
}
